/******************************************************************************
 * Copyright 2022 The Untouch Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/
#include "modules/calibration/lidar_imu_calibrator/include/odometry_frame_set.h"

#include <fstream>

#include "yaml-cpp/yaml.h"

#include "cyber/common/log.h"

namespace lidar_imu_calibrator {

OdometryFrameSet::OdometryFrameSet()
    : progress_(0), total_count_(0), complete_count_(0) {}

bool OdometryFrameSet::StartLoad(const std::string& pcd_folder_path,
                                 const std::string& odometry_file_path,
                                 const std::string& extrinsic_file_path) {
  std::vector<double> pcd_times;

  if (!LoadExtrinsic(extrinsic_file_path)) {
    return false;
  }

  if (!LoadTimestamps(pcd_folder_path + "/pcd_timestamp.csv", pcd_times)) {
    return false;
  }

  if (!LoadPoses(odometry_file_path)) {
    return false;
  }

  ResetProgress(static_cast<int>(pcd_times.size()));
  uint32_t thread_num = std::thread::hardware_concurrency();
  thread_num = std::min(thread_num, static_cast<uint32_t>(pcd_times.size()));
  AINFO << "total thread count " << thread_num;
  thread_pool_.reset(new apollo::localization::msf::ThreadPool());
  thread_pool_->start(thread_num);

  for (size_t i = 0; i < pcd_times.size(); ++i) {
    thread_pool_->commit(std::bind(&OdometryFrameSet::LoadPointCloud, this,
                                   pcd_times[i], pcd_folder_path));
  }

  return true;
}

void OdometryFrameSet::StartCompensateData() {
  ResetProgress(static_cast<int>(frames_.size()));
  uint32_t thread_num = std::thread::hardware_concurrency();
  thread_num = std::min(thread_num, static_cast<uint32_t>(frames_.size()));
  AINFO << "total thread count " << thread_num;
  thread_pool_.reset(new apollo::localization::msf::ThreadPool());
  thread_pool_->start(thread_num);

  for (size_t i = 0; i < frames_.size(); ++i) {
    thread_pool_->commit(
        std::bind(&OdometryFrameSet::CompensatePointCloud, this, frames_[i]));
  }
}

void OdometryFrameSet::Stop() {
  if (thread_pool_) {
    thread_pool_->stop();
    thread_pool_.reset();
  }
}

bool OdometryFrameSet::IsNeedCompensate() { return (frames_.size() > 0); }

const Eigen::Isometry3d& OdometryFrameSet::GetExtrinsicParameter() {
  return extrinsic_;
}

void OdometryFrameSet::SelectKeyFrames(float key_frame_delta_x,
                                       float key_frame_delta_angle) {
  if (frames_.empty()) {
    return;
  }

  key_frames_.clear();
  key_frames_.push_back(frames_.front());
  for (const auto& frame : frames_) {
    const auto& last_keyframe_pose = key_frames_.back()->GetFramePose();
    const auto& current_frame_pose = frame->GetFramePose();

    Eigen::Isometry3d delta = last_keyframe_pose.inverse() * current_frame_pose;
    double delta_x = delta.translation().norm();
    double delta_angle = Eigen::AngleAxisd(delta.linear()).angle();

    if (delta_x > key_frame_delta_x || delta_angle > key_frame_delta_angle) {
      key_frames_.push_back(frame);
    }
  }
}

void OdometryFrameSet::UpdateDelayTime(double delay_time) {
  if (key_frames_.empty()) {
    return;
  }

  for (auto& frame : key_frames_) {
    frame->SetFrameImuPose(
        InterpolatePose(frame->GetFrameTimestamp() + delay_time));
  }
}

void OdometryFrameSet::UpdateExtrinsics(const Eigen::Isometry3d& extrinsics) {
  extrinsic_ = extrinsics;
  if (key_frames_.empty()) {
    return;
  }

  for (auto& frame : key_frames_) {
    frame->UpdateFramePose(extrinsics);
  }
}

void OdometryFrameSet::DrawPointCloud(glk::GLSLShader& shader,
                                      float downsample_resolution) {
  for (const auto& key_frame : key_frames_) {
    key_frame->DrawPointCloud(shader, downsample_resolution);
  }
}

bool OdometryFrameSet::SaveExtrinsicParameter(
    const Eigen::Isometry3d& extrinsic, const std::string& file_path) {
  YAML::Node config;
  config["child_frame_id"] = child_frame_id_;
  Eigen::Quaterniond qua(extrinsic.linear());
  config["transform"]["rotation"]["x"] = qua.x();
  config["transform"]["rotation"]["y"] = qua.y();
  config["transform"]["rotation"]["z"] = qua.z();
  config["transform"]["rotation"]["w"] = qua.w();

  config["transform"]["translation"]["x"] = extrinsic.translation().x();
  config["transform"]["translation"]["y"] = extrinsic.translation().y();
  config["transform"]["translation"]["z"] = extrinsic.translation().z();

  config["header"]["frame_id"] = frame_id_;
  config["header"]["stamp"]["secs"] = 0;
  config["header"]["stamp"]["nsecs"] = 0;
  config["header"]["seq"] = 0;

  std::ofstream fout(file_path, std::ofstream::out);
  if (!fout.is_open()) {
    AERROR << "failed to open save extrinsic file!";
    return false;
  }
  fout << config;
  fout.close();
  extrinsic_ = extrinsic;
  return true;
}

void OdometryFrameSet::AddOdometryFrame(OdometryFrame::Ptr& frame) {
  std::lock_guard<std::mutex> locker(frames_mutex_);
  frames_.push_back(frame);
}

Eigen::Isometry3d OdometryFrameSet::InterpolatePose(double timestamp) {
  size_t index = std::lower_bound(ins_timestamps_.begin(),
                                  ins_timestamps_.end(), timestamp) -
                 ins_timestamps_.begin();
  index = std::max(1lu, std::min(index, ins_timestamps_.size() - 1));

  double cur_timestamp = ins_timestamps_[index];
  double pre_timestamp = ins_timestamps_[index - 1];

  double t = (cur_timestamp - timestamp) / (cur_timestamp - pre_timestamp);

  Eigen::Isometry3d pre_pose = ins_poses_[index - 1];
  Eigen::Isometry3d cur_pose = ins_poses_[index];
  Eigen::Quaterniond pre_quatd(pre_pose.linear());
  Eigen::Translation3d pre_transd(pre_pose.translation());
  Eigen::Quaterniond cur_quatd(cur_pose.linear());
  Eigen::Translation3d cur_transd(cur_pose.translation());

  Eigen::Quaterniond res_quatd = pre_quatd.slerp(1 - t, cur_quatd);

  Eigen::Translation3d re_transd;
  re_transd.x() = pre_transd.x() * t + cur_transd.x() * (1 - t);
  re_transd.y() = pre_transd.y() * t + cur_transd.y() * (1 - t);
  re_transd.z() = pre_transd.z() * t + cur_transd.z() * (1 - t);

  Eigen::Isometry3d pose = re_transd * res_quatd;

  return pose;
}

void OdometryFrameSet::LoadPointCloud(double timestamp,
                                      const std::string& pcd_folder_path) {
  const auto& pose = InterpolatePose(timestamp);
  std::string cloud_filename =
      pcd_folder_path + "/" + TimeToString(timestamp) + ".pcd";
  auto frame = std::make_shared<OdometryFrame>(cloud_filename, pose, timestamp);
  frame->LoadPointCloud(ins_timestamps_, ins_poses_, extrinsic_);

  AddOdometryFrame(frame);
  UpdateCompleteStatus();
}

void OdometryFrameSet::CompensatePointCloud(
    std::shared_ptr<OdometryFrame>& frame) {
  frame->LoadPointCloud(ins_timestamps_, ins_poses_, extrinsic_);
  UpdateCompleteStatus();
}

void OdometryFrameSet::ResetProgress(int total_count) {
  progress_ = 0;
  complete_count_ = 0;
  total_count_ = total_count;
}

bool OdometryFrameSet::LoadTimestamps(const std::string& file_path,
                                      std::vector<double>& timestamps) {
  timestamps.clear();

  FILE* file = fopen(file_path.c_str(), "r");
  if (file) {
    double timestamp;
    static constexpr int kSize = 1;
    char buffer[128] = {0};

    if (fgets(buffer, sizeof(buffer), file) != nullptr) {
      while (fscanf(file, "%lf\n", &timestamp) == kSize) {
        timestamps.push_back(timestamp);
      }
    } else {
      AERROR << "Can't read file title: " << file_path;
      return false;
    }
    fclose(file);
  } else {
    AERROR << "Can't open file to read: " << file_path;
    return false;
  }
  return true;
}

bool OdometryFrameSet::LoadPoses(const std::string& file_path) {
  ins_poses_.clear();
  ins_timestamps_.clear();

  FILE* file_ins = fopen(file_path.c_str(), "r");
  if (file_ins) {
    double timestamp;
    double x, y, z;
    double qx, qy, qz, qr;
    double std_x, std_y, std_z;
    static constexpr int kSize = 11;
    char buffer[255] = {0};
    if (fgets(buffer, sizeof(buffer), file_ins) != nullptr) {
      while (fscanf(file_ins, "%lf,%lf,%lf,%lf,%lf,%lf,%lf,%lf,%lf,%lf,%lf\n",
                    &timestamp, &x, &y, &z, &qx, &qy, &qz, &qr, &std_x, &std_y,
                    &std_z) == kSize) {
        static Eigen::Vector3d trans_init = Eigen::Vector3d(x, y, z);
        Eigen::Translation3d trans(Eigen::Vector3d(x, y, z) - trans_init);
        Eigen::Quaterniond quat(qr, qx, qy, qz);
        ins_poses_.push_back(trans * quat);
        ins_timestamps_.push_back(timestamp);
      }
    } else {
      AERROR << "Can't read file title: " << file_path;
      return false;
    }
    fclose(file_ins);
  } else {
    AERROR << "Can't open file to read: " << file_path;
    return false;
  }
  return true;
}

std::string OdometryFrameSet::TimeToString(double time, int precision) {
  std::stringstream ss;
  ss << std::fixed;
  ss.precision(precision);
  ss << time;
  return ss.str();
}

bool OdometryFrameSet::LoadExtrinsic(const std::string& file_path) {
  YAML::Node config = YAML::LoadFile(file_path);
  if (config["header"] && config["header"]["frame_id"] &&
      config["child_frame_id"] && config["transform"] &&
      config["transform"]["translation"] && config["transform"]["rotation"]) {
    frame_id_ = config["header"]["frame_id"].as<std::string>();
    child_frame_id_ = config["child_frame_id"].as<std::string>();
    extrinsic_.translation()(0) =
        config["transform"]["translation"]["x"].as<double>();
    extrinsic_.translation()(1) =
        config["transform"]["translation"]["y"].as<double>();
    extrinsic_.translation()(2) =
        config["transform"]["translation"]["z"].as<double>();

    double qx = config["transform"]["rotation"]["x"].as<double>();
    double qy = config["transform"]["rotation"]["y"].as<double>();
    double qz = config["transform"]["rotation"]["z"].as<double>();
    double qw = config["transform"]["rotation"]["w"].as<double>();
    extrinsic_.linear() = Eigen::Quaterniond(qw, qx, qy, qz).toRotationMatrix();
    return true;
  }
  AERROR << "The format of extrinsic file is wrong!";
  return false;
}

void OdometryFrameSet::UpdateCompleteStatus() {
  if (total_count_ != 0) {
    std::lock_guard<std::mutex> locker(progress_mutex_);
    ++complete_count_;
    auto progress = static_cast<int>(complete_count_ * 100.0 / total_count_);
    if (progress != progress_) {
      progress_ = progress;
      emit ProgressChanged(progress_);
    }
  }
}

void OdometryFrameSet::GetFrames(std::vector<OdometryFrame::Ptr> &frames){
  frames.clear();
  frames.insert(frames.end(), frames_.begin(), frames_.end());
}
}  // namespace lidar_imu_calibrator